import json
import streamingjson
from fastapi import WebSocket
from typing import Union
from openai.types.chat import ChatCompletionChunk
from ..schemas import ResponseMessage, MessageChunk
from .base import  SendStrategy, Sender

from ollama._types import ChatResponse, Message
class TextStrategy(SendStrategy):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
    async def send(self, data: str):
        await self.websocket.send_text(data)

class ChatCompletionChunkStrategy(SendStrategy):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
    async def send(self, data: ChatCompletionChunk):
        await self.websocket.send_json(data.model_dump_json())

class MessageChunkStrategy(SendStrategy):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
    async def send(self, data: MessageChunk):
        await self.websocket.send_json(data.model_dump_json(exclude_none=True))
class ResponseMessageStrategy(SendStrategy):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
    async def send(self, data: ResponseMessage):
        await self.websocket.send_json(data)

class ChatResponseStrategy(SendStrategy):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
    async def send(self, data: ChatResponse):
        await self.websocket.send_json(data)
class WebSocketSender(Sender):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
        self._strategies = {
            str: TextStrategy(websocket),
            ResponseMessage: ResponseMessageStrategy(websocket),
            ChatCompletionChunk: ChatCompletionChunkStrategy(websocket),
            ChatResponse: ChatResponseStrategy(websocket),
            MessageChunk: MessageChunkStrategy(websocket),
        }
    async def send(self, data):
        """根据数据类型选择策略"""
        strategy = self._strategies.get(type(data))
        if not strategy:
            raise TypeError(f"Unsupported type: {type(data)}")
        await strategy.send(data)
